import os

import torch
import torchvision

from dataset import Crack

class MyDeeplab(torch.nn.Module):
    def __init__(self):
        super().__init__()
        weights = torchvision.models.segmentation.DeepLabV3_MobileNet_V3_Large_Weights.DEFAULT
        self.deeplab = torchvision.models.segmentation.deeplabv3_mobilenet_v3_large(weights = weights)
        self.deeplab.classifier[4] = torch.nn.Conv2d(in_channels=256, out_channels=2, kernel_size=(1,1), stride=(1,1))
        self.deeplab.aux_classifier[4] = torch.nn.Conv2d(in_channels=10, out_channels=2, kernel_size=(1,1), stride=(1,1))

    def forward(self, imgs):
        result = self.deeplab(imgs)
        return result


if __name__ == '__main__':
    model = MyDeeplab()
    # print(model)
    dataset = Crack(r'./data/train/imgs', r'./data/train/masks')
    dataset_loader = torch.utils.data.DataLoader(dataset=dataset, batch_size=4, shuffle=True)

    model.eval()
    for imgs, labels in dataset_loader:
        print('imgs.shape: ', imgs.shape)
        result = model(imgs)
        print(result['out'].shape)
        break